import random
import time
import uuid
import pandas as pd

from application.db import get_database_connection
from application.db.mysql_db.info import ResourceInformationList, ResourceInformationTagsRelation, \
    ResourceInformationAttachmentList, ResourceInformationSectionList

from application.tasks.base_task import BaseTask
from application.utils.utils import generate_random_id


class DataStorageTask(BaseTask):
    @get_database_connection().atomic()  # 保证事务
    def run(self, df: pd.DataFrame):
        """
           对输入的 DataFrame 进行结构化处理，生成信息相关的多张表。

           :param df: 输入的原始 DataFrame
           :return: 包含信息表、标签表、附件表、章节表的字典
           """
        df['_id'] = df.apply(lambda _: generate_random_id("info"), axis=1)
        # 信息表
        info_rename_map = {
            "_id": "information_id",
            "information_name": "information_name",
            "page_url": "original_link",
            "marc_code": "original_language",
            "metadata": "metadata",
            "source_id": "source_id",
            "info_date": "publish_date"
        }
        # 标签关系表
        tag_rename_map = {
            "_id": "information_id",
            "tag_code": "tag_code",
            "column_info": "tag_value",
        }

        information_df = (
            df[["_id", "information_name", "page_url", "marc_code", "info_date", "metadata", "source_id"]]
            .reset_index(drop=True)
            .rename(columns=info_rename_map)
            .dropna(how="all")
        ).dropna(axis=1, how='all')

        tag_relation_df = (
            df[["_id", "tag_code", "column_info"]]
            .reset_index(drop=True)
            .rename(columns=tag_rename_map)
            .dropna(how="all")
        ).dropna(axis=1, how='all')

        # 附件表
        attachment_df = (
            df.apply(self.build_attachment_records, axis=1)
            .explode()
            .apply(pd.Series)
            .dropna(how="all")
        ).dropna(axis=1, how='all')

        # 章节表
        section_df = (
            df.apply(self.build_section_records, axis=1)
            .explode()
            .apply(pd.Series)
            .dropna(how="all")
        ).dropna(axis=1, how='all')

        # 如果 information_df 不为空才入库
        if not information_df.empty:
            ResourceInformationList.insert_many(
                self.clean_dataframe(information_df).to_dict(orient="records")
            ).execute()

        # 如果 tag_relation_df 不为空才入库
        if not tag_relation_df.empty:
            ResourceInformationTagsRelation.insert_many(
                self.clean_dataframe(tag_relation_df).to_dict(orient="records")
            ).execute()

        # 如果 attachment_df 不为空才入库
        if not attachment_df.empty:
            ResourceInformationAttachmentList.insert_many(
                self.clean_dataframe(attachment_df).to_dict(orient="records")
            ).execute()

        # 如果 section_df 不为空才入库
        if not section_df.empty:
            ResourceInformationSectionList.insert_many(
                self.clean_dataframe(section_df).to_dict(orient="records")
            ).execute()

        return df

    @staticmethod
    def build_attachment_records(row: pd.Series) -> list[dict]:
        """
        构建附件记录。

        :param row: DataFrame 的一行数据
        :return: 附件记录的字典列表
        """
        attachments = row.get("link_data", [])
        if not isinstance(attachments, list):
            attachments = []
        info_id = row.get("_id")

        return [
            {
                "information_id": info_id,
                "attachment_id": generate_random_id("atta"),
                "attachment_name": attachment.get("accessory_name"),
                "attachment_address": attachment.get("accessory_url"),
                "display_order": index + 1,
            }
            for index, attachment in enumerate(attachments)
        ]

    @staticmethod
    def build_section_records(row: pd.Series) -> list[dict]:
        """
        构建章节记录。

        :param row: DataFrame 的一行数据
        :return: 章节记录的字典列表
        """
        sections = row.get("info_section", [])
        if not isinstance(sections, list):
            sections = []
        info_id = row.get("_id")

        return [
            {"information_id": info_id,
             "section_id": generate_random_id("source"),
             "section_attr": section.get("section_attr"),
             "section_order": index + 1,
             "title_level": section.get("title_level", 0) or 0,
             "marc_code": section.get("marc_code"),
             "src_text": section.get("text_info"),
             "dst_text": section.get("dst_text"),
             "media_info": section.get("media_info"),
             }
            for index, section in enumerate(sections)
        ]
